import os
from PIL import Image
from paddleocr import PaddleOCR
import cv2
import numpy as np


# 全局配置变量 - 可根据需要修改
FOLDER_PATH = "./NEW"  # 待处理图片文件夹路径
LEFT_RATIO = 0.0   # 裁剪区域左边界比例 (0.0-1.0)
TOP_RATIO = 0.86   # 裁剪区域上边界比例 (0.0-1.0)
RIGHT_RATIO = 1.0  # 裁剪区域右边界比例 (0.0-1.0)
BOTTOM_RATIO = 0.92 # 裁剪区域下边界比例 (0.0-1.0)

# 全局变量（程序运行中自动设置）
crop_coords = None  # 存储裁剪的绝对值坐标 (left, top, right, bottom)
# 初始化OCR模型（删除不支持的det_algorithm参数，保留核心兼容参数）
ocr = PaddleOCR(
    use_textline_orientation=True,  # 启用文字方向检测（纠正倾斜文字）
    lang="ch",                      # 明确指定中文识别（关键，避免识别成英文）
    use_doc_orientation_classify=False, 
    use_doc_unwarping=False, 
    text_recognition_batch_size=1   # 单批次识别，提高稳定性（兼容新版本）
)


def preview_crop_area_on_original(file_path):
    """在原图上标记裁剪区域并显示，使用原图的百分比刻度"""
    global crop_coords
    
    try:
        img = Image.open(file_path)
        width, height = img.size
        img_np = np.array(img)
    except Exception as e:
        print(f"打开图片失败: {e}")
        return False

    # 计算裁剪区域的绝对值坐标（基于原图比例）
    left = int(LEFT_RATIO * width)
    top = int(TOP_RATIO * height)
    right = int(RIGHT_RATIO * width)
    bottom = int(BOTTOM_RATIO * height)
    crop_coords = (left, top, right, bottom)

    # 绘制裁剪区域矩形（绿色框，醒目）
    cv2.rectangle(img_np, (left, top), (right, bottom), (0, 255, 0), 2)
    
    # 绘制横向百分比刻度（红色线+文字）
    for i in range(11):
        x = int(width * i / 10)
        cv2.line(img_np, (x, 0), (x, height), (0, 0, 255), 1)
        cv2.putText(
            img_np, f"{i*10}%", (x + 5, 30),
            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2
        )

    # 绘制纵向百分比刻度（红色线+文字）
    for i in range(11):
        y = int(height * i / 10)
        cv2.line(img_np, (0, y), (width, y), (0, 0, 255), 1)
        cv2.putText(
            img_np, f"{i*10}%", (10, y + 15),
            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2
        )

    # 显示裁剪区域比例和操作提示
    crop_info = f"裁剪区域: 左={LEFT_RATIO*100}%, 上={TOP_RATIO*100}%, 右={RIGHT_RATIO*100}%, 下={BOTTOM_RATIO*100}%"
    cv2.putText(img_np, crop_info, (20, height - 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2)
    cv2.putText(img_np, "按 'y' 确认并继续, 按 'ESC' 退出", (20, height - 60), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 0, 0), 2)

    # 调整窗口大小（避免图片太大超出屏幕）
    max_display_size = 1000
    scale = min(max_display_size / width, max_display_size / height, 1.0)
    if scale < 1.0:
        img_np = cv2.resize(img_np, (int(width * scale), int(height * scale)))

    # 显示预览窗口
    cv2.imshow('原图裁剪区域预览（带刻度）', img_np)
    key = cv2.waitKey(0)
    cv2.destroyAllWindows()

    return key == ord('y')  # 按'y'确认，其他键取消


def recognize_text_from_image(file_path):
    """对图片的指定区域进行裁剪并识别文字"""
    global crop_coords
    if not crop_coords:
        print("裁剪坐标未初始化")
        return ""

    try:
        # 打开图片并裁剪指定区域
        img = Image.open(file_path)
        cropped_img = img.crop(crop_coords)
        
        # 保存裁剪图到临时目录（方便调试：查看裁剪是否准确）
        temp_dir = "./temp_crops"
        os.makedirs(temp_dir, exist_ok=True)
        temp_path = os.path.join(temp_dir, f"temp_{os.path.basename(file_path)}")
        cropped_img.save(temp_path)
        
        # 转换为OpenCV兼容的BGR格式（PaddleOCR两种格式都支持，保险起见统一）
        img_np = np.array(cropped_img)
        if len(img_np.shape) == 3 and img_np.shape[2] == 3:
            img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
        
        # 执行OCR识别（核心：调用ocr.ocr()，无多余参数）
        result = ocr.predict(img_np)
        
        # 打印原始结果（调试用：看识别到的原始数据）
        print(f"图片 {os.path.basename(file_path)} 的原始识别结果：{result}")
        
        # 提取文字（适配PaddleOCR返回格式）
        text = ""
        text = result[0].get("rec_texts")[0]
        
        return text.strip()  # 去除首尾空格/换行
    except Exception as e:
        print(f"处理图片 {os.path.basename(file_path)} 时出错: {str(e)}")
        return ""


def batch_rename_images():
    """批量处理：预览第一张→确认后批量裁剪+识别+重命名"""
    # 筛选文件夹中的图片文件
    image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.gif')
    image_files = [f for f in os.listdir(FOLDER_PATH) if f.lower().endswith(image_extensions)]

    if not image_files:
        print("指定文件夹中没有找到图片文件")
        return

    print(f"找到 {len(image_files)} 个图片文件，准备处理...")

    # 预览第一张图片，确认裁剪区域
    first_file = os.path.join(FOLDER_PATH, image_files[0])
    print(f"正在预览第一张图片: {image_files[0]}")
    if not preview_crop_area_on_original(first_file):
        print("用户取消操作，程序终止")
        return

    # 批量处理所有图片
    for file in image_files:
        file_path = os.path.join(FOLDER_PATH, file)
        print(f"\n处理文件: {file}")

        # 识别文字
        recognized_text = recognize_text_from_image(file_path)
        if not recognized_text:
            print("未能识别到文字，跳过该文件")
            continue

        # 处理文件名（去除非法字符，避免重命名失败）
        invalid_chars = '/\\:*?"<>|'
        new_file_name = ''.join([c for c in recognized_text if c not in invalid_chars])
        file_ext = os.path.splitext(file)[1]  # 保留原始扩展名（如.jpg）
        new_file_path = os.path.join(FOLDER_PATH, new_file_name + file_ext)

        # 避免文件名重复（重复则加后缀_1、_2...）
        counter = 1
        while os.path.exists(new_file_path):
            new_file_path = os.path.join(FOLDER_PATH, f"{new_file_name}_{counter}{file_ext}")
            counter += 1

        # 执行重命名
        try:
            os.rename(file_path, new_file_path)
            print(f"重命名成功: {file} -> {new_file_name}{file_ext}")
        except Exception as e:
            print(f"重命名失败: {e}")


if __name__ == '__main__':
    batch_rename_images()
    print("\n所有文件处理完毕")